/*
 * CudaAesOperationsTest.cpp
 *
 *  Created on: Feb 21, 2010
 *      Author: Łukasz Spintzyk
 */
//#define DBG_NO_CUDA
#ifndef DBG_NO_CUDA
#include "CudaAESOperations.h"
#endif
#include "aesUtils.h"

#include "AesKeyTestVectors.h"
#include <iostream>
#include <iomanip>

#include <cppunit/TestRunner.h>
#include <cppunit/TestResult.h>
#include <cppunit/TestResultCollector.h>
#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/BriefTestProgressListener.h>
#include <cppunit/extensions/TestFactoryRegistry.h>

class TestRoundTransformations: public CPPUNIT_NS::TestCase
{
CPPUNIT_TEST_SUITE(TestRoundTransformations);
		//CPPUNIT_TEST(testSubByte4x4_shared);
		CPPUNIT_TEST(testShiftRow4x4);
		CPPUNIT_TEST(testShiftRow4x6);
		CPPUNIT_TEST(testShiftRow4x8);
		CPPUNIT_TEST(testMixColumn4x4);
		CPPUNIT_TEST(testMixColumn4x4);
		CPPUNIT_TEST(testMixColumn4x6);
		CPPUNIT_TEST(testMixColumn4x8);
		CPPUNIT_TEST(testAddRoundKey4x4);
		CPPUNIT_TEST(testSubByte4x4);
		CPPUNIT_TEST(testSubByte4x4_2);
		CPPUNIT_TEST(testSubByte4x4_3);
		CPPUNIT_TEST(testSubByte4x6);
		CPPUNIT_TEST(testSubByte4x8);



		CPPUNIT_TEST(testAddRoundKey4x4);
		CPPUNIT_TEST(testAddRoundKey4x6);
		CPPUNIT_TEST(testAddRoundKey4x8);
		CPPUNIT_TEST(testAesRound);
		CPPUNIT_TEST(testAesBlockEncryption);
	CPPUNIT_TEST_SUITE_END();

public:
	void setUp(void)
	{
		d_blockSize = 4*8;
		d_roundKeySize = (NR[NB_8][NB_8]+1)*4*BLOCK_SIZE(NB_8)/4;
		d_block = new byte[d_blockSize];
		d_copyBlock = new byte[d_blockSize];
		d_roundKey = new byte[d_roundKeySize];

		expandKey((byte*)c_key256, NB_4, d_roundKey, NB_4, 0);

		setTestBlock(d_block, 8, 4);
		//setTestBlock(d_roundKey,d_roundKeySize/4, 4, 152);
		memcpy(d_copyBlock, d_block, sizeof(byte) * d_blockSize);

		aesEncryptInit();
	}
	void tearDown(void)
	{
		for(int i = 0; i < d_blockSize; ++i)
		{
			d_block[i] = d_copyBlock[i] = 0;
		}

		delete[] d_block;
		delete[] d_copyBlock;
		delete[] d_roundKey;
	}

protected:
	void testAesBlockEncryption(void)
	{
		const unsigned int size = 4 * 4;
		//expandKey(c_key192FF,NB_4, d_e);
		aesEncryptBlock(d_block, d_roundKey, NB_4, NB_4);
		rijndaelTest(d_copyBlock, d_roundKey, NB_4, NB_4);

		//std::cout << "\nRound key"<<std::endl;
		//printBlock(d_roundKey, d_roundKeySize/12,4);

		//std::cout << "key len " <<d_roundKeySize <<std::endl;

		//std::cout << "\nAfter rijndael"<<std::endl;
		//printBlock(d_block, size/4,4);
		//std::cout << "\nAfter rijndael2"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}


	void testAesRound(void)
	{
		const unsigned int size = 4 * 4;
		//expandKey(c_key192FF,NB_4, d_e);
		aesRound(d_block, d_roundKey, NB_4);

		byteSub(d_copyBlock, NB_4);
		shiftRow(d_copyBlock, NB_4);
		mixColumn(d_copyBlock, NB_4);
		addRoundKey(d_copyBlock, d_roundKey, NB_4);

		//std::cout << "\nAfter round"<<std::endl;
		//printBlock(d_block, size/4,4);
		//std::cout << "\nAfter round2"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}

	void testSubByte4x4(void)
	{
		const unsigned int size = 4 * 4;
		d_copyBlock[3] = d_block[3] = 0xFB;

		byteSub(d_block, NB_4);


		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX[d_copyBlock[i]]);
		}
	}

	void testSubByte4x4_shared(void)
	{
		const unsigned int size = 4 * 4;
		d_copyBlock[3] = d_block[3] = 0xFB;
		//std::cout << "\nbefore bytesub"<<std::endl;
		//printBlock(d_block, size/4,4);
		byteSub(d_block, NB_4);
		//std::cout << "\nAfter bytesub"<<std::endl;
		//printBlock(d_block, size/4,4);

		//clock_t start = clock();


		//std::cout << "\nAfter bytesub2"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);
		//clock_t end = clock();
		//std::cout << "CPU Byte sub time = "<< end-start << std::endl;

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX[d_copyBlock[i]]);
		}
	}

	void testSubByte4x4_2(void)
	{
		const unsigned int size = 4 * 4;
		byte data[size] = {
				0x3e, 0x24, 0x43, 0xfb, 0x74, 0xea, 0xdc, 0x6d, 0x39, 0x82, 0x9c, 0xe4, 0x9c, 0x4b, 0x52, 0xae
		};

		byte expected[16]=
		{
			0xb2, 0x36, 0x1a, 0x0f, 0x92, 0x87, 0x86, 0x3c, 0x12, 0x13, 0xde, 0x69, 0xde, 0xb3, 0x00, 0xe4
		};


		memcpy(d_block, data, 16*sizeof(byte));

		byteSub(d_block, NB_4);


		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX[data[i]]);
			CPPUNIT_ASSERT_EQUAL(d_block[i], expected[i]);
		}
	}

	void testSubByte4x4_3(void)
	{
		const unsigned int size = 4 * 4;
		byte* blockIn = new byte[256];
		byte* blockOut = new byte[256];
		for (int i =0; i < 256; ++i)
		{
			memset(d_block,(byte) i, 16);
			memset(d_copyBlock,(byte) i, 16);

			byteSub(d_block, NB_4);

			for (int j = 0; j < size; j++)
			{
				CPPUNIT_ASSERT_EQUAL(d_block[j], SBOX[d_copyBlock[j]]);
			}
		}

	}


	void testSubByte4x6(void)
	{
		const unsigned int size = 4 * 6;
		byteSub(d_block, NB_6);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX[d_copyBlock[i]]);
		}
	}

	void testSubByte4x8(void)
	{
		const unsigned int size = 4 * 8;
		byteSub(d_block, NB_8);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], SBOX[d_copyBlock[i]]);
		}
	}

	void testShiftRow4x4(void)
	{
		const unsigned int size = 4 * 4;
		//std::cout << "\nbefore shift"<<std::endl;
		//printBlock(d_block, size/4,4);
		shiftRow(d_block, NB_4);
		//std::cout << "\nAfter shift"<<std::endl;
		//printBlock(d_block, size/4,4);
		shiftRow2(d_copyBlock, NB_4);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}

	void testShiftRow4x6(void)
	{
		const unsigned int size = 4 * 6;
		shiftRow(d_block, NB_6);
		shiftRow2(d_copyBlock, NB_6);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}

	void testShiftRow4x8(void)
	{
		const int size = 4 * 8;
		//std::cout << "\nBefore shift"<<std::endl;
		//printBlock(d_block, size/4,4);
		shiftRow(d_block, NB_8);
		shiftRow2(d_copyBlock, NB_8);
		//std::cout << "\nAfter shift"<<std::endl;
		//printBlock(d_block, size/4,4);
		//std::cout << "\nAfter2 shift"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);
		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}
	void testMixColumn4x4(void)
	{
		const int size = 4 * 4;
									   //after mixcolumn
		d_copyBlock[0]=d_block[0]=0xd4;//04
		d_copyBlock[1]=d_block[1]=0xbf;//66
		d_copyBlock[2]=d_block[2]=0x5d;//81
		d_copyBlock[3]=d_block[3]=0x30;//e5
		//std::cout << "\nBefore Mix coumn"<<std::endl;
		//printBlock(d_block, size/4,4);

		mixColumn(d_block, NB_4);

		//std::cout << "\nAfter Mix coumn"<<std::endl;
		//printBlock(d_block, size/4,4);

		mixColumnTest(d_copyBlock, 4, 4, (byte*) MIXTAB, 4, 4);

		//std::cout << "\nAfter Expecterd Mix coumn"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);

		CPPUNIT_ASSERT(d_block[0]==0x4);
		CPPUNIT_ASSERT(d_block[1]==0x66);
		CPPUNIT_ASSERT(d_block[2]==0x81);
		CPPUNIT_ASSERT(d_block[3]==0xe5);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}
	void testMixColumn4x6(void)
	{
		const int size = 4 * 6;
									   //after mixcolumn
		d_copyBlock[0]=d_block[0]=0xd4;//04
		d_copyBlock[1]=d_block[1]=0xbf;//66
		d_copyBlock[2]=d_block[2]=0x5d;//81
		d_copyBlock[3]=d_block[3]=0x30;//e5

		mixColumn(d_block, NB_6);
		mixColumnTest(d_copyBlock, 6, 4, (byte*) MIXTAB, 4, 4);

		CPPUNIT_ASSERT(d_block[0]==0x4);
		CPPUNIT_ASSERT(d_block[1]==0x66);
		CPPUNIT_ASSERT(d_block[2]==0x81);
		CPPUNIT_ASSERT(d_block[3]==0xe5);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}
	void testMixColumn4x8(void)
	{
		//expandKey<NB_8> (d_block, 32);
		//std::cout << "expand "<< d_block[0] <<std::endl;
		const int size = 4 * 8;
									   //after mixcolumn
		d_copyBlock[0]=d_block[0]=0xd4;//04
		d_copyBlock[1]=d_block[1]=0xbf;//66
		d_copyBlock[2]=d_block[2]=0x5d;//81
		d_copyBlock[3]=d_block[3]=0x30;//e5

		mixColumn(d_block, NB_8);
		mixColumnTest(d_copyBlock, 8, 4, (byte*) MIXTAB, 4, 4);

		//std::cout << "\nAfter shift"<<std::endl;
		//printBlock(d_block, size/4,4);
		//std::cout << "\nAfter copymix"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);

		CPPUNIT_ASSERT(d_block[0]==0x4);
		CPPUNIT_ASSERT(d_block[1]==0x66);
		CPPUNIT_ASSERT(d_block[2]==0x81);
		CPPUNIT_ASSERT(d_block[3]==0xe5);

		CPPUNIT_ASSERT(memcmp(d_block, d_copyBlock, size) == 0);
	}

	void testAddRoundKey4x4(void)
	{
		const unsigned int size = 4 * 4;
		//std::cout << "\n Before addround"<<std::endl;
		//printBlock(d_block, size/4,4);
		addRoundKey(d_block, d_roundKey, NB_4);

		addRoundKeyTest(d_copyBlock, d_roundKey, NB_4);
		//std::cout << "\nAfter addround"<<std::endl;
		//printBlock(d_block, size/4,4);
		//std::cout << "\nAfter addround2"<<std::endl;
		//printBlock(d_copyBlock, size/4,4);
		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], d_copyBlock[i]);
		}
	}

	void testAddRoundKey4x6(void)
	{
		const unsigned int size = 4 * 6;
#ifndef DBG_NO_CUDA
		addRoundKey(d_block, d_roundKey, NB_6);
#endif
		addRoundKeyTest(d_copyBlock, d_roundKey, NB_6);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], d_copyBlock[i]);
		}
	}

	void testAddRoundKey4x8(void)
	{
		const unsigned int size = 4 * 8;
#ifndef DBG_NO_CUDA
		addRoundKey(d_block, d_roundKey, NB_8);
#endif
		addRoundKeyTest(d_copyBlock, d_roundKey, NB_8);

		for (int i = 0; i < size; i++)
		{
			CPPUNIT_ASSERT_EQUAL(d_block[i], d_copyBlock[i]);
		}
	}

private:
	void rijndaelTest(byte* state,byte* d_key, NB nb, NB nk)
	{
		unsigned int bw = BLOCK_SIZE(nb)/4;

		addRoundKey(state, d_key, NB_4);
		for(int i =1; i < NR[nk][nb]; ++i)
		{
			aesRound(state, d_key + 4*i*bw, NB_4);
		}

		//final round
		byteSub(state, NB_4);
		shiftRow(state, NB_4);
		addRoundKey(state, d_key + 4*bw*NR[nk][nb], NB_4);

	}

	void shiftRow2(byte* state, NB nb)
	{
		unsigned int size = BLOCK_SIZE(nb);
		byte* state_copy = new byte[size * sizeof(byte)];
		memcpy(state_copy, state, size);

		int w = size / 4;
		for (int x = 0; x < w; ++x)
		{
			int x2 = (x + w - C1(nb)) % w;
			state_copy[x2 * 4 + 1] = state[x * 4 + 1];
			assert(x2*4+1 < size);
			assert(x*4+1 < size);

			x2 = (x + w - C2(nb)) % w;
			state_copy[x2 * 4 + 2] = state[x * 4 + 2];
			x2 = (x + w - C3(nb)) % w;
			state_copy[x2 * 4 + 3] = state[x * 4 + 3];
		}
		memcpy(state, state_copy, size);
		delete[] state_copy;
	}

	void setTestBlock(byte* block, int w, int h, int seed=1)
	{
		int k = seed;
		for (int i = 0; i < w; ++i)
		{
			for (int j = 0; j < h; ++j)
			{
				block[i * h + j] = k++;
			}
		}
	}

	void printBlock(byte* block, int w, int h)
	{
		for (int j = 0; j < h; ++j)
		{
			for (int i = 0; i < w; ++i)
			{
				std::cout << std::hex << std::setw(4) << (unsigned int) block[i * h + j]
						<< " ";
			}
			std::cout << std::endl;
		}
	}

	void mixColumnTest(
			byte* state, int w1, int h1,
			byte* CBlock, int w2, int h2)
	{

		for (int i = 0; i < w1; ++i)
		{
			byte* column = new byte[h1];
			for (int j = 0; j < h1; ++j)
			{
				byte* p = state+i*h1+j;
				column[j] = *p;
			}
			assert(h1==4);
			assert(w2 == 4 );
			assert(h2 == 4);
			matMul(column, h1, CBlock, w2, h2);
			for (int j = 0; j < h1; ++j)
			{
				byte* p = state+i*h1+j;
				*p = column[j];
			}
			delete[] column;
		}
	}

	byte field28Mno(byte c1, byte c2)
	{
		if (c2 == 1)
		{
			return c1;
		}
		else if(c2 == 2)
		{
			return XTIME(c1);
		}
		else if (c2 ==3)
		{
			return XTIME(c1)^c1;
		}
	}

	void matMul(byte* a, int ha, byte* b, int wb, int hb)
	{
		byte* copy = new byte[ha];

		assert(ha == hb);

		for (int i = 0; i < ha; ++i)
		{
			byte sum = 0;
			sum =0;
			for (int j = 0; j < wb; ++j)
			{
				byte* p = b+i*ha +j;
				sum ^= field28Mno(a[j], *p);//*p * a[j];
			}
			copy[i] = sum;
		}

		memcpy(a, copy, ha);

		delete[] copy;
	}

	void addRoundKeyTest(byte* state, byte* key,const NB nb)
	{
	    unsigned int size = BLOCK_SIZE(nb);

	    for(int i =0; i < size; ++i)
	    {
	    	state[i]^=key[i];
	    }
	}

	byte* d_block;
	byte* d_copyBlock;
	byte* d_roundKey;
	int d_blockSize;
	int d_roundKeySize;

};
CPPUNIT_TEST_SUITE_REGISTRATION(TestRoundTransformations);

int main(int ac, char **av)
{
	//--- Create the event manager and test controller
	CPPUNIT_NS::TestResult controller;
	//--- Add a listener that colllects test result
	CPPUNIT_NS::TestResultCollector result;
	controller.addListener(&result);
	//--- Add a listener that print dots as test run.
	CPPUNIT_NS::BriefTestProgressListener progress;
	controller.addListener(&progress);
	//--- Add the top suite to the test runner
	CPPUNIT_NS::TestRunner runner;
	runner.addTest(CPPUNIT_NS::TestFactoryRegistry::getRegistry().makeTest());
	runner.run(controller);

	return result.wasSuccessful() ? 0 : 1;
}
